import numpy as np
import os
import pandas as pd
from libpmf.python import libpmf
import vecs_io


def delete_dir_if_exist(dire):
    if os.path.isdir(dire):
        command = 'rm -rf %s' % dire
        print(command)
        os.system(command)


if __name__ == '__main__':
    # filename_l = ['movielens-small', 'movielens-1m', 'movielens-10m', 'movielens-27m']
    # filename_l = ['amazon', 'book-crossing', 'goodreads', 'movielens-27m', 'netflix', 'yahoomusic', 'yahoomusic_big']
    # filename_l = ['amazon-home-kitchen', 'movielens-27m', 'netflix', 'yahoomusic', 'yahoomusic_big', 'yelp']
    # filename_l = ['amazon-home-kitchen', 'yahoomusic_big', 'movielens-1b', 'yelp']
    # filename_l = ['ml-1m']
    filename_l = ['lastfm']
    # n_dim = 150
    for n_dim in [150]:
        for filename in filename_l:
            base_dir = './intermediate-rating-csv'

            csv = pd.read_csv(os.path.join(base_dir, '%s.csv' % filename))
            n_user = np.max(csv['userID'])
            n_item = np.max(csv['itemID'])
            res = libpmf.train_coo(row_idx=csv['userID'] - 1, col_idx=csv['itemID'] - 1, obs_val=csv['rating'],
                                   m=n_user,
                                   n=n_item, param_str='-k %d -n 20' % n_dim)
            user_l = res['W']
            item_l = res['H']

            save_dir = '/home/bianzheng/Dataset/MIPS'
            save_filename = os.path.join(save_dir, '%s-%dd' % (filename, n_dim))
            delete_dir_if_exist(save_filename)
            os.mkdir(save_filename)

            save_user = os.path.join(save_filename, '%s_user.fvecs' % filename)
            save_item = os.path.join(save_filename, '%s_item.fvecs' % filename)
            vecs_io.fvecs_write(save_user, user_l)
            vecs_io.fvecs_write(save_item, item_l)
            print(filename, n_dim, 'complete n_user {}, n_item {}'.format(n_user, n_item))
